package com.ppch.wuwamanus.agent;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.ppch.wuwamanus.advisors.DatabaseChatMemory;
import com.ppch.wuwamanus.agent.model.AgentState;
import com.ppch.wuwamanus.agent.model.ReasonData;
import com.ppch.wuwamanus.config.CosClientConfig;
import com.ppch.wuwamanus.config.FileUtils;
import com.ppch.wuwamanus.mapper.ChatHistoryMapper;
import com.ppch.wuwamanus.model.entity.ChatHistory;
import com.ppch.wuwamanus.model.enums.ChatHistoryMessageTypeEnum;
import com.ppch.wuwamanus.service.ChatHistoryService;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.schema.Database;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

/**
 * Package:ppch-ai-agent
 * ClassName:ToolCallAgent
 *
 * @Author 泡泡茶壶
 * @Create 2025/7/30 21:23
 * @Version 1.0
 * Description:
 * 工具调用智能体：定义具体如何调用工具完成任务
 * 1.因为Spring AI 已经将先思考再行动的模式实现并封装了，所以我们无法直接实现 reason() 和 act()方法。
 * 解决方案 1：只实现 reason()、act()中的其中一个方法，该方案实现简单、快速
 * 解决方案 2：取消Spring AI 的自动调用工具机制，自定义控制工具执行生命周期 官方文档：https://docs.spring.io/spring-ai/reference/api/tools.html#_user_controlled_tool_execution
 * 但是由于我们使用的是阿里云百炼的 DashScope 灵积模型API服务，不能通过官方的方法来取消，
 * 需要使用DashScopeChatOptions.builder().withProxyToolCalls(true).build();来取消Spring AI的自动调用工具机制
 * 解决方案 3：自己实现工具调用能力：自己写 Prompt，引导 AI 回复想要调用的工具列表和调用参数，然后再执行工具并将结果
 * 返回给 AI 再次执行。（复杂且效果可能不如官方实现）
 * 2.下面计划采用方案 2， 自定义控制工具执行生命周期
 * 3.自定义控制工具调用实现参考官网：https://docs.spring.io/spring-ai/reference/api/tools.html#_user_controlled_tool_execution
 */
@EqualsAndHashCode(callSuper = true)
@Slf4j
@Data
@Component
public class ToolCallAgent extends ReActAgent {

    /**
     * 工具调用管理器
     */
    private ToolCallingManager toolCallingManager;

    /**
     * 可以调用的工具集合
     */
    private ToolCallback[] availableToolCalls;

    /**
     * MCP 服务
     */
    private ToolCallbackProvider toolCallbackProvider;

    /**
     * 聊天模型选项
     */
    private ChatOptions chatOptions;

    /**
     * AI思考结果
     */
    private ChatResponse chatResponse;

    /**
     * 文件工具类（将文件上传到 COS 对象存储中）
     */
    private FileUtils fileUtils;

    /**
     * 对话历史 Mapper
     */
    private ChatHistoryMapper chatHistoryMapper;

    /**
     * 对话历史 ChatHistoryService
     */
    private ChatHistoryService chatHistoryService;

    /**
     * 当前用户 ID
     */
    private Long userId;

    /**
     * COS 客户端配置
     */
    private CosClientConfig cosClientConfig;


    /**
     * 无参构造
     */
    public ToolCallAgent() {
    }

    /**
     * 初始化
     */
    public ToolCallAgent(ToolCallback[] availableToolCalls, ToolCallbackProvider toolCallbackProvider,
                         FileUtils fileUtils,ChatHistoryMapper chatHistoryMapper,CosClientConfig cosClientConfig,
                         ChatHistoryService chatHistoryService,Long userId
    ) {
        super();
        this.availableToolCalls = availableToolCalls;
        this.toolCallbackProvider = toolCallbackProvider;
        this.toolCallingManager = ToolCallingManager.builder().build();
        //TODO:取消Spring AI的自动调用工具机制
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
        this.fileUtils = fileUtils;
        this.chatHistoryMapper = chatHistoryMapper;
        this.cosClientConfig = cosClientConfig;
        this.chatHistoryService = chatHistoryService;
        this.userId = userId;
    }

    /**
     * 加载聊天记录到 List<Message> memoryList 中，实现多轮对话功能
     * 1.只加载最新的 chatMemoryRetrieveSize 条对话记录
     * 2.按照对话记录的创建时间升序排序（从旧到新）
     * @param chatId 对话 ID
     * @param chatMemoryRetrieveSize 关联的上下文大小
     * @return 数据条数
     */
    public int loadChatHistoryToMemory(Long chatId, int chatMemoryRetrieveSize) {
        try {
            //为了避免 memoryList 加载重复数据，首先将 memoryList 清空
            this.getMemoryList().clear();
            LambdaQueryWrapper<ChatHistory> queryWrapper = new LambdaQueryWrapper<>();
            queryWrapper.eq(ChatHistory::getChatId, chatId);
            // 改为按时间降序排列，确保通过limit获取最新的N条记录
            queryWrapper.orderByDesc(ChatHistory::getCreateTime);
            queryWrapper.last("LIMIT " + chatMemoryRetrieveSize);
            List<ChatHistory> chatHistories = chatHistoryMapper.selectList(queryWrapper);
            if (chatHistories == null || chatHistories.isEmpty()) {
                return 0;
            }
            // 将结果反转，使消息按时间升序排列（从旧到新）
            Collections.reverse(chatHistories);
            //已成功加载的数据条数
            int loadedCount = 0;
            //将聊天记录加载到 memoryList 中
            for (ChatHistory chatHistory : chatHistories) {
                Message message = null;
                String messageType = chatHistory.getMessageType();
                String chatHistoryMessage = chatHistory.getMessage();
                if (ChatHistoryMessageTypeEnum.USER.getValue().equals(messageType)) {
                    //添加用户消息
                    message = new UserMessage(chatHistoryMessage);
                } else if (ChatHistoryMessageTypeEnum.AI.getValue().equals(messageType) ||
                        ChatHistoryMessageTypeEnum.SEND_PDF_FILE.getValue().equals(messageType) ||
                        ChatHistoryMessageTypeEnum.SEND_CARD_DRAWING_ANALYSIS.getValue().equals(messageType)) {
                    //AI消息、发送PDF文件消息、发送抽卡分析数据消息都属于AssistantMessage，都需要添加
                    message = new AssistantMessage(chatHistoryMessage);
                }
                if (message != null) {
                    this.getMemoryList().add(message);
                    loadedCount++;
                }
            }
            log.info("成功为 chatId 为：" + chatId + "的对话加载了" + loadedCount + "条聊天记录");
            return loadedCount;
        } catch (Exception e) {
            log.error("从数据库加载聊天记录失败，chatId：{}, error：{}", chatId, e.getMessage());
            return 0;
        }
    }

    /**
     * 调用一次 AI获取响应结果，从中得到要调用的工具列表，并判断是否需要调用工具（包含 AI 的思考过程）
     *
     * @return 是否需要调用工具 true：需要 false：不需要
     */
    @Override
    public ReasonData reason() {
        //TODO：0.从数据库中加载历史对话，实现多轮对话
        //loadChatHistoryToMemory(this.getChatId(), this.getChatMemoryRetrieveSize());
        //1.校验下一步操作提示词，作为用户提示词并保存到上下文中
        if (StrUtil.isNotBlank(this.getNextStepPrompt()) && this.getNextStepPrompt() != null) {
            UserMessage userMessage = new UserMessage(this.getNextStepPrompt());
            this.getMemoryList().add(userMessage);
        }
        try {
            //2.将用户提示词发送给AI，得到工具调用的响应结果
            List<Message> memoryList = this.getMemoryList();
            Prompt prompt = new Prompt(memoryList, chatOptions);
            ChatResponse chatResponse = this.getChatClient()
                    .prompt(prompt)
                    .system(this.getSystemPrompt())
                    //将自定义工具提供给 ChatClient
                    .tools(availableToolCalls)
                    //将 MCP 服务中的所有工具提供给 ChatClient
                    .tools(toolCallbackProvider)
                    .call()
                    .chatResponse();
            //工具响应结果
            this.chatResponse = chatResponse;
            //3.解析结果，获取要调用的工具列表
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //AI决定要调用的工具列表
            List<AssistantMessage.ToolCall> toBeInvokedToolList = assistantMessage.getToolCalls();
            //大模型调用结果
            String assistantMessageText = assistantMessage.getText();
            //TODO：打印日志（AI 思考结果，也就是AI使用原始能力回复的结果）
            String aiReasonResult = StrUtil.isEmpty(assistantMessageText) ? "" : assistantMessageText;
            log.info("ToolCallAgent reason() 打印日志，AI 的思考结果为："+aiReasonResult);
            //AI选择调用的工具数量
            String invokedToolsNumber = toBeInvokedToolList.size() == 0 ? "" : "需要调用" + toBeInvokedToolList.size() + "个工具，";
            //AI选择调用的工具列表：包括工具名称和工具参数
            String toolCallString = toBeInvokedToolList.stream().map(toolCall ->
                    String.format("工具名称：%s，工具参数：%s。", toolCall.name(), toolCall.arguments())
            ).collect(Collectors.joining("\n"));
            //组合文案
            String toolCallResult = aiReasonResult + invokedToolsNumber + toolCallString;
            //TODO：打印工具调用日志
            log.info("reason()方法中打印日志，AI选择调用工具，结果为：" + toolCallResult);
            //4.如果要调用的工具列表非空并且文案也为不需要调用工具，则返回false
            if (toBeInvokedToolList.isEmpty() && !toolCallResult.contains("需要调用1个工具")) {
                //无需调用工具则将助手消息（AI的思考结果）添加到上下文（无需调用工具返回false）
                this.getMemoryList().add(assistantMessage);
                return new ReasonData(false, toolCallResult);
            } else {
                //需要调用工具，则将AI选择的工具信息添加到上下文（需要调用工具返回true）
                return new ReasonData(true, toolCallResult);
            }
        } catch (Exception e) {
            //5.异常处理
            log.error("reason()方法 246行 打印日志：" + this.getName() + "思考过程出现错误：" + e.getMessage());
            //this.getMemoryList().add(new AssistantMessage("出现异常：" + e.getMessage()));
            return new ReasonData(false, "");
        }
    }

    /**
     * 调用工具完成对应任务，并返回调用结果。
     *
     * @return 工具调用结果
     */
    @Override
    public String act() {
        //1.获取最新的用户提示词
        Prompt prompt = new Prompt(this.getMemoryList(), chatOptions);
        //2.调用工具
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, this.chatResponse);
        //3.获取工具调用后的消息历史记录，包括助手消息和工具执行结果。
        List<Message> conversationHistory = toolExecutionResult.conversationHistory();
        //TODO: 4.更新消息上下文（很重要，如果没有更新，AI 就会重复执行第一步，因为上下文中一直没有第一步的执行结果）
        this.setMemoryList(conversationHistory);

        //获取工具响应信息对象
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollectionUtil.getLast(conversationHistory);
        //获取工具调用结果列表
        List<ToolResponseMessage.ToolResponse> responses = toolResponseMessage.getResponses();
        //5.判断是否执行了终止工具 TerminateTool 中的 doTerminate()方法，如果执行了则改变智能体的执行状态为Finished
        boolean doTerminated = responses
                .stream()
                .anyMatch(toolResponse -> toolResponse.name().equals("doTerminate"));
        if (doTerminated) {
            this.setState(AgentState.FINISHED);
        }
        //TODO：将工具调用结果保存到 Map<String, String> toolCallResult 中，便于后续获取各个工具的调用结果
        //      比如根据generatePDF工具名称获取工具结果，将结果作为事件发送给前端。
        for (ToolResponseMessage.ToolResponse toolResponse : responses) {
            Map<String, String> toolCallResult = this.getToolCallResult();
            String toolName = toolResponse.name();
            String toolResult = toolResponse.responseData();
            toolCallResult.put(toolName, toolResult);
        }

        //TODO：使用 DatabaseChatMemory 实现对话记忆持久化（需要手动将工具结果保存到数据库中）
        //      因为现在是自己控制工具调用的，Spring AI 无法帮我们将工具调用结果保存到 DatabaseChatMemory 中
        DatabaseChatMemory databaseChatMemory = new DatabaseChatMemory(chatHistoryService);
        databaseChatMemory.setCurrentUserId(this.getUserId());
        List<Message> messages = new ArrayList<>();
        messages.add(new ToolResponseMessage(responses));
        databaseChatMemory.add(this.getChatId().toString(),messages);


        //获取工具调用结果（定义输出格式），并将结果返回
        return toolResponseMessage.getResponses().stream()
                .map(toolResponse -> "工具执行结果如下: " + toolResponse.responseData()).collect(Collectors.joining("\n"));
    }

    /**
     * 智能体的多步骤执行 Agent Loop，接收用户Prompt，通过多步骤执行生成最终回复。
     *
     * @param userPrompt 用户 Prompt
     * @return Flux<String>
     */
    public Flux<String> runReactiveStream(String userPrompt) {
        Flux<String> flux = Flux.create(
                sink -> {
            //1.用户提示词为空，不能执行
            try {
                if (StrUtil.isBlank(userPrompt)) {
                    sink.next("用户提示词为空");
                    //完成传输
                    sink.complete();
                    return;
                }
                //2.如果智能体不是空闲状态，不能执行
                if (this.getState() != AgentState.IDLE) {
                    sink.next("当前状态不能执行智能体：" + this.getState());
                    //完成传输
                    sink.complete();
                    return;
                }
            } catch (Exception e) {
                log.error("runReactiveStream 482 行执行异常：" + e.getMessage());
                //sink.error(e);
                return;
            }
            //3.执行状态改为执行中 RUNNING
            this.setState(AgentState.RUNNING);
            //4.将用户提示词记录到消息上下文
            this.getMemoryList().add(new UserMessage(userPrompt));
            //5.多步骤执行（循环），try catch防止执行失败
            try {
                //执行步骤数不能超过最大步骤且执行状态不能为Finished
                for (int i = 0; i < this.getMaxSteps() && this.getState() != AgentState.FINISHED; i++) {
                    this.setCurrentSteps(i + 1);
                    log.info("已执行至{}/{}步", this.getCurrentSteps(), this.getMaxSteps());
                    //调用单步执行方法获取对应的执行结果
                    String currentStepResult = this.step();
                    String toLogResult = "step" + this.getCurrentSteps() + "结果为:" + (StrUtil.isEmpty(currentStepResult) ? "无结果" : currentStepResult);
                    log.info("ToolCallAgent类中的runReactiveStream()方法打印日志（469行）：" + toLogResult);
                    //添加到执行结果列表
                    List<String> resultList = this.getResultList();
                    resultList.add(currentStepResult);
                    //TODO：如果当前执行结果与上一步执行结果相同，则跳过当前执行结果，避免发送重复数据
                    // 结果集合必须大于 1，避免把第一次的响应数据也给置为空了。
                    if(resultList.size() > 1){
                        String preResult = resultList.get(resultList.size() - 2);
                        if(StrUtil.equals(currentStepResult,preResult)){
                            log.warn("将重复响应置空，被置空数据为：{}",preResult);
                            sink.next("");
                        }
                    }else {
                        //将每一步的执行结果发送出去
                        sink.next(currentStepResult);
                    }

                    //TODO：监听工具调用结果，如果生成了 PDF 文件，则将文件 URL 发送给前端
                    Map<String, String> toolCallResult = this.getToolCallResult();
                    if(toolCallResult.containsKey("generatePDF") && MapUtil.isNotEmpty(toolCallResult)){
                        //获取 PDF 生成工具的结果：PDF 文件保存到项目本地的 URL路径）
                        String pdfUrl = toolCallResult.get("generatePDF");
                        //获取文件名
                        String pdfFileName = fileUtils.getPdfFileName(pdfUrl);
                        log.info("pdfFileName:" + pdfFileName);
                        //TODO: 将文件上传到 COS 对象存储
                        fileUtils.uploadFileToCOS(pdfFileName);
                        // 发送自定义事件(发送PDF文件)
                        String data = cosClientConfig.getHost() +"/file/" + pdfFileName;
                        sink.next(ChatHistoryMessageTypeEnum.SEND_PDF_FILE.getValue() + data);
                    }
                    //TODO：监听工具调用结果，发送 SEND_CARD_DRAWING_ANALYSIS 事件，携带抽卡分析结果数据
                    if(toolCallResult.containsKey("getCardDrawingAnalysisResult") && MapUtil.isNotEmpty(toolCallResult)){
                        // 发送自定义事件
                        sink.next(ChatHistoryMessageTypeEnum.SEND_CARD_DRAWING_ANALYSIS.getValue()
                                + toolCallResult.get("getCardDrawingAnalysisResult"));
                    }
                }
                //6.当执行步骤数达到最大，将执行状态修改为Finished
                if (this.getCurrentSteps() >= this.getMaxSteps()) {
                    this.setState(AgentState.FINISHED);
                    //this.getResultList().add("执行结束，已达最大执行步骤数：" + this.getMaxSteps());
                    //sink.next("执行结束，已达最大执行步骤数：" + this.getMaxSteps());
                    log.warn("runReactiveStream 529 行 执行结束，已达最大执行步骤数：{}",this.getMaxSteps());
                }
                //完成传输
                sink.complete();
            } catch (Exception e) {
                //8.执行异常
                this.setState(AgentState.ERROR);
                log.error("runReactiveStream 536 行 执行异常：" + e.getMessage(), e); // 添加完整的异常堆栈记录
                //sink.error(e);
            }
        });
        return flux;
    }


    //region
    /**
     * （同步）智能体的多步骤执行 Agent Loop，接收用户 Prompt，通过多步骤执行生成最终回复
     *
     * @param userPrompt 用户 Prompt
     * @return 最终的执行结果
     */
    public String run(String userPrompt) {
        return "该方法具体实现请看 BaseAgent 类";
    }

    /**
     * （ SSE 流式输出）智能体的多步骤执行 Agent Loop，接收用户 Prompt，通过多步骤执行生成最终回复。
     * TODO：由于该方式不便于在 Controller 层收集 AI 响应数据，从而将响应数据存入数据库中，所以使用下面的 runReactiveStream 代替）
     * @param userPrompt 用户 Prompt
     * @return SseEmitter
     */
    public SseEmitter runStream(String userPrompt) {
        log.info("用户提示词为：" + userPrompt);
        //创建 SseEmitter 对象，可以设置超时时间
        SseEmitter sseEmitter = new SseEmitter(1800000L); // 30 min
        //TODO：将下列代码改为异步执行，避免阻塞主线程，且允许在每个步骤(step())完成后立即发送结果，实现真正的流式输出。
        CompletableFuture.runAsync(() -> {
            //1.用户提示词为空，不能执行
            try {
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("用户提示词为空");
                    //完成传输
                    sseEmitter.complete();
                    return;
                }
                //2.如果智能体不是空闲状态，不能执行
                if (this.getState() != AgentState.IDLE) {
                    sseEmitter.send("当前状态不能执行智能体：" + this.getState());
                    //完成传输
                    sseEmitter.complete();
                    return;
                }
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
            //3.执行状态改为执行中 RUNNING
            this.setState(AgentState.RUNNING);
            //4.将用户提示词记录到消息上下文
            this.getMemoryList().add(new UserMessage(userPrompt));
            //5.多步骤执行（循环），try catch防止执行失败
            try {
                //执行步骤数不能超过最大步骤且执行状态不能为Finished
                for (int i = 0; i < this.getMaxSteps() && this.getState() != AgentState.FINISHED; i++) {
                    this.setCurrentSteps(i + 1);
                    log.info("已执行至{}/{}步", this.getCurrentSteps(), this.getMaxSteps());
                    //调用单步执行方法获取对应的执行结果
                    String stepResult = this.step();
                    String result = "step" + this.getCurrentSteps() + ":" + stepResult;
                    log.info(result);
                    //添加到执行结果列表
                    this.getResultList().add(result);
                    //将每一步的执行结果发送出去
                    sseEmitter.send(result);
                    //TODO：监听工具调用结果，如果生成了 PDF 文件，则将文件 URL 发送给前端
                    Map<String, String> toolCallResult = this.getToolCallResult();
                    if(toolCallResult.containsKey("generatePDF") && MapUtil.isNotEmpty(toolCallResult)){
                        //获取 PDF 生成工具的结果：PDF 文件保存到项目本地的 URL路径）
                        String pdfUrl = toolCallResult.get("generatePDF");
                        //获取文件名
                        String pdfFileName = fileUtils.getPdfFileName(pdfUrl);
                        log.info("pdfFileName:" + pdfFileName);
                        //将文件上传到COS对象存储
                        fileUtils.uploadFileToCOS(pdfFileName);
                        // 发送自定义事件(发送PDF文件)
                        SseEmitter.SseEventBuilder event = SseEmitter.event()
                                .name(ChatHistoryMessageTypeEnum.SEND_PDF_FILE.getValue())
                                //将文件保存到对象存储中，便于前端用户访问和下载
                                .data(cosClientConfig.getHost() +"/file/" + pdfFileName);
                        sseEmitter.send(event);
                    }
                    //TODO：监听工具调用结果，发送 SEND_CARD_DRAWING_ANALYSIS 事件，携带抽卡分析结果数据
                    if(toolCallResult.containsKey("getCardDrawingAnalysisResult") && MapUtil.isNotEmpty(toolCallResult)){
                        // 发送自定义事件
                        SseEmitter.SseEventBuilder event = SseEmitter.event()
                                .name(ChatHistoryMessageTypeEnum.SEND_CARD_DRAWING_ANALYSIS.getValue())
                                .data(toolCallResult.get("getCardDrawingAnalysisResult"));
                        sseEmitter.send(event);
                    }
                }
                //6.当执行步骤数达到最大，将执行状态修改为Finished
                if (this.getCurrentSteps() >= this.getMaxSteps()) {
                    this.setState(AgentState.FINISHED);
                    this.getResultList().add("执行结束，已达最大执行步骤数：" + this.getMaxSteps());
                    sseEmitter.send("执行结束，已达最大执行步骤数：" + this.getMaxSteps());
                }
                //完成传输
                sseEmitter.complete();
            } catch (Exception e) {
                //8.执行异常
                this.setState(AgentState.ERROR);
                log.error("执行异常：" + e.getMessage());
                try {
                    sseEmitter.send("执行异常" + e.getMessage());
                    sseEmitter.complete();
                } catch (IOException ex) {
                    //完成传输
                    sseEmitter.completeWithError(e);
                }
            }
        });

        //连接超时回调函数
        sseEmitter.onTimeout(() -> {
            this.setState(AgentState.ERROR);
            log.error(" SseEmitter 连接超时");
        });

        //完成回调函数
        sseEmitter.onCompletion(() -> {
            if (this.getState() == AgentState.RUNNING) {
                this.setState(AgentState.FINISHED);
                log.info(" SseEmitter 任务完成");
            }
        });
        return sseEmitter;
    }
    //endregion
}
